# gradio_app.py
import gradio as gr
import paddle
import numpy as np
from PIL import Image
import paddlenlp as ppnlp

# 模型定义
class MultiModalClassifier(paddle.nn.Layer):
    def __init__(self, num_classes):
        super().__init__()
        self.image_encoder = paddle.vision.models.resnet50(pretrained=True)
        self.image_encoder.fc = nn.Identity()
        self.image_proj = paddle.nn.Linear(2048, 512)
        
        self.text_encoder = ppnlp.transformers.ErnieModel.from_pretrained('ernie-1.0')
        self.text_proj = paddle.nn.Linear(768, 512)
        
        self.fusion = paddle.nn.Sequential(
            paddle.nn.Linear(1024, 512),
            paddle.nn.ReLU(),
            paddle.nn.Dropout(0.5),
            paddle.nn.Linear(512, 256),
            paddle.nn.ReLU(),
            paddle.nn.Dropout(0.5)
        )
        
        self.classifier = paddle.nn.Linear(256, num_classes)
    
    def forward(self, image, input_ids, attention_mask):
        image_features = self.image_encoder(image)
        image_features = self.image_proj(image_features)
        
        text_outputs = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask)
        text_features = text_outputs[1]
        text_features = self.text_proj(text_features)
        
        fused_features = paddle.concat([image_features, text_features], axis=1)
        fused_features = self.fusion(fused_features)
        logits = self.classifier(fused_features)
        return logits

# 加载模型
model = MultiModalClassifier(num_classes=5)
model.set_state_dict(paddle.load("checkpoints/best_model.pdparams"))
model.eval()

# 加载tokenizer
tokenizer = ppnlp.transformers.ErnieTokenizer.from_pretrained('ernie-1.0')

# 定义类别映射
label_map = {
    0: "科技",
    1: "娱乐",
    2: "体育",
    3: "财经",
    4: "教育"
}

# 图像预处理
def preprocess_image(image):
    image = image.resize((224, 224))
    image = np.array(image).astype('float32')
    image = image / 255.0
    image = (image - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225])
    image = np.transpose(image, (2, 0, 1))
    return paddle.to_tensor(image, dtype='float32').unsqueeze(0)

# 文本预处理
def preprocess_text(text, max_seq_len=128):
    encoded_inputs = tokenizer(
        text=text,
        max_seq_len=max_seq_len,
        pad_to_max_seq_len=True,
        return_attention_mask=True,
        return_token_type_ids=False
    )
    input_ids = paddle.to_tensor(encoded_inputs['input_ids'], dtype='int64').unsqueeze(0)
    attention_mask = paddle.to_tensor(encoded_inputs['attention_mask'], dtype='int64').unsqueeze(0)
    return input_ids, attention_mask

# 预测函数
def predict(image, text):
    if image is None or text == "":
        return "请上传图像并输入文本描述"
    
    # 处理图像
    image_tensor = preprocess_image(image)
    
    # 处理文本
    input_ids, attention_mask = preprocess_text(text)
    
    # 模型推理
    with paddle.no_grad():
        logits = model(image_tensor, input_ids, attention_mask)
        probs = paddle.nn.functional.softmax(logits, axis=1).numpy()[0]
    
    # 构建结果
    results = {label_map[i]: float(probs[i]) for i in range(len(probs))}
    
    return results

# 创建Gradio界面
def create_interface():
    with gr.Blocks(title="多模态分类系统") as interface:
        gr.Markdown("# 多模态分类系统")
        gr.Markdown("上传一张图片并输入相关文本描述，系统将自动分类")
        
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="上传图片")
                text_input = gr.Textbox(label="输入文本描述")
                predict_btn = gr.Button("分类预测")
            
            with gr.Column():
                output = gr.Label(label="分类结果")
        
        predict_btn.click(
            fn=predict,
            inputs=[image_input, text_input],
            outputs=output
        )
        
        # 添加示例
        gr.Examples(
            examples=[
                ["examples/tech.jpg", "这款新手机的相机功能非常出色"],
                ["examples/sports.jpg", "这支足球队在本赛季表现出色"],
                ["examples/entertainment.jpg", "这部电影的特效令人震撼"]
            ],
            inputs=[image_input, text_input]
        )
    
    return interface

# 启动界面
if __name__ == "__main__":
    interface = create_interface()
    interface.launch(server_name="0.0.0.0", server_port=7860)